import pcap, libnet
import sys, os, termios, fcntl, threading, copy, thread
import socket, string, struct, argparse, time
import logging, exceptions, random, errno
from libnet.constants import *
from datetime import datetime
from collections import deque
from PyQt4 import QtCore, QtGui
from packetParser import packetParser

from fuzzers import *

protocols = {socket.IPPROTO_UDP:'udp',
             socket.IPPROTO_TCP:'tcp'}

global rules, proxyList, proxyList2, logSettings, spoofingOpts
proxyList = []
proxyList2 = []
rules = {}
logSettings = {}
spoofingOpts = [False]

proxyQueue = deque()
loggingQueue = deque()
packetQueue = deque()

Qlock = thread.allocate_lock()
proxyLock = thread.allocate_lock()

decoder = packetParser()

global receivedPackages,forwarderPackages,fuzzedPackages,duplicatedPackages,outputPackages
receivedPackages = 0
forwardedPackages = 0
fuzzedPackages = 0
duplicatedPackages = 0
outputPackages = 0

    
def setRules(proxy_state, flt, log, spoofing=False):
    try:
        global rules, spoofingOpts
        
        rules['proxy'] = proxy_state
        rules.update(flt)

        if rules['level']!=0:
            decoder.setLevel(rules['level'])

        #TODO: Logging settings
        #print log

        if spoofing:
            pass
    
    except KeyError as e:
        print "Rules were erronous: {}".format(e)


def resetVariables():
    receivedPackages = 0
    forwardedPackages = 0
    fuzzedPackages = 0
    duplicatedPackages = 0
    '''Is this wise thing to do??'''
    #loggingQueue.clear()
    #packetQueue.clear()
    

'''
Function for checking available interfaces and validity of the interfaces
given in command line arguments.
'''
def checkInterfaces(interface):

    devices = []
    if(pcap.findalldevs()):
        for name, descr, addrs, flags in pcap.findalldevs():
            if addrs:
                devices.append(name)

    else:
        return -1
    
    for i in interface:
        if i not in devices:
            return -1

    return 0


'''
This class will take care of forwarding incoming packages from output
interface to input interface, but only if proxy mode is allowed to be.
Packages are forwarded using OSI layer 2.
'''
class proxyThread(QtCore.QThread):

    def __init__(self, interface, parent=None):
        QtCore.QThread.__init__(self,parent)
        self.interface = interface

        #Socket for forwarding stuff
        self.sock = socket.socket(socket.AF_PACKET, socket.SOCK_RAW)
        self.sock.bind((self.interface,0))
    
        self.sLimit = 65535

        self.operate = True


    def run(self):

        while self.operate:
            
            #proxyLock.acquire()

            while len(proxyQueue)!=0:
                
                packet = proxyQueue.popleft()

                #proxyLock.release()

                global outputPackages
                outputPackages+=1

                global proxyList2
                proxyList2.append(packet)

                try:
                    self.sock.send(packet)
            
                except socket.error as e:
                    pLen = len(packet)
                    if cmp(e[0],errno.EMSGSIZE)==0 and pLen<self.sLimit:
                        if not modifyMTU(pLen, self.sock, self.interface[1]):
                            self.sLimit = pLen
                except:
                    print "Unhandled exception: {}".format(sys.exc_info()[0])

                packet = ""

            #else:
            #proxyLock.release()
            time.sleep(0.01)
                        

    '''
    In case of too large packet for interface to send, try to modify
    interface's MTU value by calling directly ioctl function which is
    responsible for it in unix environment.
    '''
    def modifyMTU(self, pktlen, sck, iface):
        
        print "Packet too large, trying to modify interface's MTU value."
        try:
            fcntl.ioctl(sck, 0x8922,
                        struct.pack('<16sH',iface,pktlen)+'\x00'*14)
            print "MTU value for interface '{}' changed to {}".format(iface,pktlen)
            return True
        except Exception as ex:
            print "MTU value could not be changed, ignoring packets larger than {} bytes".format(pktlen-1)
            return False
        

    def shutDown(self):

        self.operate = False


'''
This class reads incoming packages from a queue and process them by
going through rules and determining if package needs to be fuzzed
or just forwarded.
'''
class forwarderThread(QtCore.QThread):

    updateInfoDisplay = QtCore.pyqtSignal(dict, int)

    def __init__(self, interface, fuzzerlist, log, duplOpt, fOpts, parent=None):
        QtCore.QThread.__init__(self,parent)
        self.interface = str(interface)
        self.fuzzerList = copy.deepcopy(fuzzerlist)
        self.log = log

        self.p = {}
        self.sLimit = 65535
        self.operate = True
        self.sendAgainCtr = 2
        self.duplOpt = duplOpt
        self.duplOpt['activate'] = False
        self.duplArray = []

        if self.duplOpt['state']:
            for i in range(0,100): self.duplArray.append(False)
            for i in range(0,int(self.duplOpt['p']*100)):
                self.duplArray[i] = True

        self.sock = socket.socket(socket.AF_PACKET, socket.SOCK_RAW)
        self.sock.bind((self.interface,0))

        #Launcher class interacts with the fuzzers
        if fOpts:
            self.launcher = Launcher(fuzzerlist, fOpts)
        else:
            self.launcher = Launcher(fuzzerlist, [])
        #self.launcher = Launcher(self.fuzzerList, [])

        self.l = libnet.context(LINK, self.interface)
        

    '''
    Function calculates checksum for given ip header.
    First all the 16 bit words are summed, then the value is divided into 2 16 bit
    words and then they are summed. Finally it takes 1's complement and returns.
    '''
    def checksum(self, msg):

        ip_header_sum = sum([struct.unpack("!H",msg[i:i+2])[0] for i in range(0,len(msg),2)])
        ip_header_sum = (ip_header_sum & 0xffff)+(ip_header_sum >> 16 & 0xffff)
        #Python doesn't support 1's complement so it has to be done like this
        ip_header_sum = ~ip_header_sum & 0xffff
        
        return struct.pack("!H", ip_header_sum)


    def run(self):

        while self.operate:

            #Qlock.acquire() #Lock on

            while len(packetQueue)!=0:

                tempPacket = packetQueue.popleft()

                #Qlock.release()

                global receivedPackages
                receivedPackages+=1
                
                self.p['raw'] = tempPacket[0]
                self.p['length'] = tempPacket[1]

                #If fixed addresses are defined, try to change it 
                if spoofingOpts[0]:
                    #If IPv4 packet, change the address and calculate checksum
                    if (("sa" in spoofingOpts or "da" in spoofingOpts) and len(self.p['raw'])>33 and 
                        socket.ntohs(struct.unpack('H',self.p['raw'][12:14])[0])==2048):
                        try:
                            originalRaw = self.p['raw']
                            originalSA = ""
                            originalDA = ""
                            header = "{}{}{}".format(self.p['raw'][14:24],"\x00\x00",self.p['raw'][26:34])
                            if 'sa' in spoofingOpts:
                                originalSA = header[12:16]
                                addr=socket.inet_aton(spoofingOpts[spoofingOpts.index('sa')+1])
                                header = "{}{}{}".format(header[0:12],addr,header[16:20])
                                header = "{}{}{}".format(header[0:10],self.checksum(header),header[12:20])
                            if 'da' in spoofingOpts:
                                originalDa = header[16:20]
                                addr=socket.inet_aton(spoofingOpts[spoofingOpts.index('da')+1])
                                header = "{}{}".format(header[0:16],addr)
                                header = "{}{}{}".format(header[0:10],self.checksum(header),header[12:20])
                            self.p['raw']="{}{}{}".format(self.p['raw'][0:14],header,self.p['raw'][34:])                                
                        except:
                            print "Error with ip spoofing."
                            self.p['raw'] = originalRaw
                            
                    #Change ethernet addresses if so defined in spoofing options
                    if ("es" in spoofingOpts or "ed" in spoofingOpts) and len(self.p['raw'])>13:
                        try:
                            originalRaw = self.p['raw']
                            originalES = ""
                            originalED = ""
                            header = self.p['raw'][0:14]
                            if 'es' in spoofingOpts:
                                originalES = header[6:12]
                                header = "{}{}{}".format(header[0:6],spoofingOpts[spoofingOpts.index('es')+1],header[12:14])
                            if 'ed' in spoofingOpts:
                                originalED = header[0:6]
                                header = "{}{}".format(spoofingOpts[spoofingOpts.index('ed')+1],header[6:14])
                            self.p['raw']="{}{}".format(header,self.p['raw'][14:])
                        except:
                            print "Error with ethernet spoofing."
                            self.p['raw'] = originalRaw

                #Qlock.release() #Lock off

                try:
                    '''If duplication enabled, randomly decide if packet should
                    be duplicated. If so, set a flag and wait until packet is
                    duplicated before trying randomly again.'''
                    if self.duplOpt['state'] and not self.duplOpt['activate']:
                        if random.choice(self.duplArray):
                            self.duplOpt['activate'] = True
                
                    #No sense in going further if fuzzing is not allowed
                    if not self.fuzzerList:
                        self.forwarderFunction(self.p['raw'])
                        raise NoError()
            
                    temp = decoder.packetParse(self.p['raw'])
                    #Parser should return negative value if parsing were to be erronous
                    if temp<0:
                        self.forwarderFunction(self.p['raw'])
                        raise NoError()
                    self.p.update(temp)

                    if len(self.p['payload'])==0:
                        self.forwarderFunction(self.p['raw'])
                        raise NoError()

                    #This goes through the filtering rules, if any
                    
                    if rules['enabled']:

                        fValue = False

                        if rules['sa']:
                            if 'sa' in spoofingOpts and originalSA!="":
                                if cmp(socket.inet_ntoa(originalSA),rules.sa)==0:
                                    fValue = True
                                else: raise ValueError
                            else:
                                if cmp(self.p['ip_source_address'],rules.sa)==0:
                                    fValue = True
                                else: raise ValueError

                        if rules['da']:
                            if 'da' in spoofingOpts and originalDA!="":
                                if cmp(socket.inet_ntoa(originalDA),rules.da)==0:
                                    fValue = True
                                else: raise ValueError
                            else:
                                if cmp(self.p['ip_destination_address'],rules.da)==0:
                                    fValue = True
                                else: raise ValueError

                        if rules['protocols']:
                            if (protocols[self.p['ip_protocol']] in rules['protocols']):
                                fValue = True
                            else: raise ValueError

                        if rules['sp']!=0:
                            if ('udp_source_port' in self.p and
                                cmp(self.p['udp_source_port'],rules['sp'])==0):
                                fValue = True
                            elif ('tcp_source_port' in self.p and
                                  cmp(self.p['tcp_source_port'],rules['sp'])==0):
                                fValue = True
                            else: raise ValueError

                        if rules['dp']!=0:
                            if ('udp_destination_port' in self.p and
                                cmp(self.p['udp_destination_port'],rules['dp'])==0):
                                fValue = True
                            elif ('tcp_destination_port' in self.p and
                                  cmp(self.p['tcp_destination_port'],rules['dp'])==0):
                                fValue = True
                            else: raise ValueError

                        if not fValue:
                            self.forwarderFunction(self.p['raw'])
                            raise NoError()
                        

                    #Packet is accepted and payload should be fuzzed
                    self.fuzzerFunction()

                except KeyError:
                    self.forwarderFunction(self.p['raw'])
                except ValueError:
                    self.forwarderFunction(self.p['raw'])
                except NoError:
                    pass

                self.p['receivedNro'] = receivedPackages
                self.p['fuzzedNro'] = fuzzedPackages
                self.p['forwardedNro'] = forwardedPackages
                self.updateInfoDisplay.emit(copy.deepcopy(self.p), 0)
                self.p.clear()

            #else:
            #Qlock.release() #Lock off
            time.sleep(0.01)


    def fuzzerFunction(self):

        try:
            #Fuzz the payload by calling the fuzzing class
            data = self.launcher.dataFuzz(self.p['payload'],self.fuzzerList[0])

            #Packet building from the highest layer to the lowest

            #UDP header
            if self.p['udp'] and cmp(protocols[self.p['ip_protocol']],"udp")==0:
                udpHeader = self.l.build_udp(
                    sp=self.p['udp_source_port'],
                    dp=self.p['udp_destination_port'],
                    sum=0,
                    payload=data)

            #TCP header
            elif self.p['tcp'] and cmp(protocols[self.p['ip_protocol']],"tcp")==0:
                #Including data to the packet
                tcpData = self.l.build_data(payload=data)

                tcpHeader = self.l.build_tcp(
                    len=self.p['tcp_data_offset'],
                    sp=self.p['tcp_source_port'],
                    dp=self.p['tcp_destination_port'],
                    seq=self.p['tcp_seqnum'],
                    ack=self.p['tcp_acknum'],
                    control=self.p['tcp_flags'],
                    win=self.p['tcp_window_size'],
                    sum=0,
                    urg=self.p['tcp_urgent_pointer'])
                if self.p['tcp_data_offset']>5:
                    self.l.build_tcp_options(options=self.p['tcp_options'])

            #IPv4 header
            if self.p['ip'] and self.p['ip_version']==4:

                if not self.p['tcp'] and not self.p['udp']:
                    self.l.build_data(payload=data)
                        
                source=self.l.name2addr4(self.p['ip_source_address'], DONT_RESOLVE)
                destination=self.l.name2addr4(self.p['ip_destination_address'], DONT_RESOLVE)
                ipv4Header = self.l.build_ipv4(
                    tos=self.p['ip_tos'],
                    id=self.p['ip_id'],
                    frag=self.p['ip_fragment_offset'],
                    ttl=self.p['ip_ttl'],
                    prot=self.p['ip_protocol'],
                    sum=0,
                    src=source,
                    dst=destination)
                if self.p['ip_header_len']>5:
                    self.l.build_ipv4_options(option=self.p['ip_options'])

            #IPv6 header
            elif self.p['ip'] and self.p['ip_version']==6:

                if not self.p['tcp'] and not self.p['udp']:
                    self.l.build_data(payload=data)

                ipv6Header = self.l.build_ipv6(
                    tc=self.p['ip_traffic_class'],
                    fl=self.p['ip_flow_label'],
                    nh=self.p['ip_next_header'],
                    hl=self.p['ip_hop_limit'],
                    src=self.p['ip_source_address'],
                    dst=self.p['ip_destination_address'])

            #Ethernet header
            if self.p['eth']:

                if not self.p['ip']:
                    self.l.build_data(payload=data)
                            
                self.l.build_ethernet(
                    dst=self.p['eth_mac_destination'],
                    src=self.p['eth_mac_source'],
                    type=self.p['eth_type'])

            packetToSend = self.l.getpacket_raw()

            if rules['proxy']:
                global proxyList
                proxyList.append(packetToSend)

            try:
                self.sock.send(packetToSend)

                if self.log['enabled']:
                    if self.log['mode']=='pcap':
                        loggingQueue.append([repr(time.time()),packetToSend])
                    elif self.log['mode']=='default':
                        port = [0,0]
                        message=["len {},".format(len(packetToSend))]

                        if self.p['udp']:
                            port[0] = self.p['udp_source_port']
                            port[1] = self.p['udp_destination_port']
                            message.append("UDP-packet,")
                        elif self.p['tcp']:
                            port[0] = self.p['tcp_source_port']
                            port[1] = self.p['tcp_destination_port']
                            message.append("TCP-packet,")

                        if self.p['ip'] and self.p['ip_version']==4:
                            message.append("IPv4,")
                        elif self.p['ip'] and self.p['ip_version']==6:
                            message.append("IPv6,")
                            
                        message.append("{}:{} > {}:{}, fuzzer: {}".format(self.p['ip_source_address'],
                                                                   port[0],
                                                                   self.p['ip_destination_address'],
                                                                   port[1],
                                                                   self.getFuzzer()))
                        loggingQueue.append([
                            datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f"),
                            ' '.join(message),
                            packetToSend])

                if self.duplOpt['activate']:
                    self.duplOpt['activate'] = False
                    self.sock.send(packetToSend)
                    global duplicatedPackages
                    duplicatedPackages+=1
                    if self.log['enabled']:
                        if self.log['mode']=='pcap':
                            loggingQueue.append([repr(time.time()),packetToSend])
                        elif self.log['mode']=='default':
                            loggingQueue.append([
                                datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f"),
                                packetToSend])

            except socket.error as e:
                dataLen = len(packetToSend)
                if cmp(e[0],errno.EMSGSIZE)==0 and dataLen<self.sLimit:
                    if not self.modifyMTU(dataLen,self.sock,self.interface):
                        self.sLimit = dataLen

            global fuzzedPackages
            fuzzedPackages+=1
            
            self.l.clear_packet()

            self.sendAgainCtr = 2

        except KeyError as werr:
            print "KeyError while building package: {}".format(werr)
        except ValueError as eks:
            '''
            If something unexpected happens while building/sending
            a packet, try to send it again at least once.
            '''
            print "Unhandled exception: {}".format(eks)

            if self.sendAgainCtr>0:

                #Qlock.acquire() #Lock on
            
                packetQueue.appendleft(self.p['raw'])

                #Qlock.release() #Lock off

                self.sendAgainCtr-=1

            else:
                self.sendAgainCtr = 2


    def fuzzerChange(self, fList):

        self.fuzzerList = copy.deepcopy(fList)


    '''Just forwards the packet using layer 2 and checks possible errors'''
    def forwarderFunction(self, data):

        if rules['proxy']:
            global proxyList
            proxyList.append(data)

        try:
            self.sock.send(data)

            if self.log['enabled']:
                if self.log['mode']=='pcap':
                    loggingQueue.append([repr(time.time()),data])
                elif self.log['mode']=='default':
                    message="len {}, forwarded packet".format(len(data))
                    loggingQueue.append([
                        datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f"),
                        message,
                        data])

            if self.duplOpt['activate']:
                #Just here to give some delay between packets
                #time.sleep(0.001)
                if rules['proxy']: proxyList.append(data)
                self.duplOpt['activate'] = False
                self.sock.send(data)
                global duplicatedPackages
                duplicatedPackages+=1
                if self.log['enabled']:
                    if self.log['mode']=='pcap':
                        loggingQueue.append([repr(time.time()),data])
                    elif self.log['mode']=='default':
                        loggingQueue.append([
                            datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f"),
                            "Duplicated packet.",
                            data])
            
        except socket.error as e:
            dataLen = len(data)
            if cmp(e[0],errno.EMSGSIZE)==0 and dataLen<self.sLimit:
                if not self.modifyMTU(dataLen,self.sock,self.interface):
                    self.sLimit = dataLen
        except Exception:
            print "Unhandled exception: {}".format(sys.exc_info()[0])
        
        global forwardedPackages
        forwardedPackages+=1
        

    '''
    In case of too large packet for interface to send, try to modify
    interface's MTU value by calling directly ioctl function which is
    responsible for it in unix environment.
    '''
    def modifyMTU(self, pktlen, sck, iface):
        print "Packet too large, trying to modify interface's MTU value."
        try:
            fcntl.ioctl(sck, 0x8922,
                        struct.pack('<16sH',iface,pktlen)+'\x00'*14)
            print "MTU value for interface '{}' changed to {}".format(iface,pktlen)
            return True
        except Exception as ex:
            print "MTU value could not be changed, ignoring packets larger than {} bytes".format(pktlen-1)
            return False


    def shutDown(self):

        self.operate = False


class loggerThread(QtCore.QThread):

    def __init__(self, logfile, options, parent=None):
        QtCore.QThread.__init__(self.parent)
        self.operate = True
        self.logfile = logfile
        self.orderNumber = 0
        self.options = options

    def run(self):

        if cmp(self.options['mode'],'pcap')==0:

            logger = open(self.logfile, 'w')

            #Adding global header to the file
            logger.write("\xa1\xb2\xc3\xd4\x00\x02\x00\x04"+
                         "\x00\x00\x00\x00\x00\x00\x00\x00"+
                         "\x00\x00\xFF\xFF\x00\x00\x00\x01")

            conv=lambda x:''.join([chr(int(x[i:i+2],16)) for i in range(0,len(x),2)])
            pad = lambda y: y.zfill(8)
            padRight = lambda w: w+''.join(["0" for i in range(len(w),6)])
            
            while self.operate:
                while len(loggingQueue)>0:

                    packet = loggingQueue.popleft()
                
                    #Pcap packet header + actual data
                    sec, micro = packet[0].split('.')
                    if len(micro)!=6: micro = padRight(micro)
                    sec, micro = hex(int(sec))[2:], hex(int(micro))[2:]
                    micro = pad(micro)
                    sec, micro = conv(sec), conv(micro)
                    length = pad(hex(len(packet[1]))[2:])
                    length = conv(length)
                    logger.write(sec+micro+length+length+packet[1])

                time.sleep(0.1)

            logger.close()

        else:
            logging.basicConfig(filename=self.logfile,
                                format='%(message)s',
                                level=logging.DEBUG)

            logging.info(datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f")+
                        " : -- Session started --")

            while self.operate:
                while len(loggingQueue)>0:

                    packet = loggingQueue.popleft()
                    self.orderNumber += 1

                    logging.info(("#%d  %s : %s\n[%s]")%
                                (self.orderNumber,
                                packet[0],
                                packet[1],
                                ','.join(["%02X" % ord(x) for x in packet[2]])))                            

                time.sleep(0.1)

            logging.info(datetime.now().strftime("%Y-%m-%d %H:%M:%S,%f")+
                         " : -- Session ended --")
                

    def shutDown(self):

        self.operate = False


class Listener(QtCore.QThread):

    updateSignal = QtCore.pyqtSignal()

    def __init__(self, interfaces, parent=None):
        QtCore.QThread.__init__(self, parent)
        self.interfaces = map(str,interfaces)

        self.operate = True
        self.paused = False


    def run(self):

        pcapObj = pcap.pcapObject()

        maxBytes = 65535
        promiscuous = False
        readTimeout = 100

        pcapObj.open_live(self.interfaces[0], maxBytes, promiscuous, readTimeout)
        pcapObj.setnonblock(1)

        if rules['proxy']:
            pcapObj2 = pcap.pcapObject()
            pcapObj2.open_live(self.interfaces[1], maxBytes, promiscuous, readTimeout)
            pcapObj2.setnonblock(1)

            self.proxyHandler = proxyThread(self.interfaces[0])
            self.proxyHandler.start()

            while self.operate:

                pcapObj.dispatch(1, self.packetHandler)

                pcapObj2.dispatch(1, self.proxyPacketHandler)

            self.proxyHandler.shutDown()

        else:
            while self.operate:

                pcapObj.dispatch(1, self.packetHandler)
            

    def packetHandler(self, pktlen, data, timestamp):

        if not data:
            return
        else:
            if rules['proxy']:
                global proxyList2
                try:
                    index = proxyList2.index(data)
                    proxyList2.pop(index)
                    if index>0: proxyList.pop(index-1)
                    return
                
                except ValueError:
                    pass

            Qlock.acquire()
            
            packetQueue.append([data,pktlen])

            Qlock.release()

    def proxyPacketHandler(self, pktlen, data, timestamp):

        if not data:
            return
        else:
            try:
                global proxyList
                index = proxyList.index(data)
                proxyList.pop(index)
                if index>0: proxyList.pop(index-1)
                return
            
            except ValueError:
                pass

            proxyLock.acquire()
            
            proxyQueue.append(data)

            proxyLock.release()

    #Doesn't work simply by implementing a while loop!
    def hooverPause(self, value):
        
        self.paused = value
    

    def shutDown(self):
        
        self.operate = False


'''Not actually an error, just a way to skip things without return'''
class NoError(Exception):
    def __init__(self):
        pass
